# app.py
from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
import os
import torch
from PIL import Image
import io
import numpy as np
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from model import CustomNet

app = Flask(__name__)

# 配置
UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CustomNet()
model.load_state_dict(torch.load('./models/final_model.pth', map_location=device))
model.to(device)
model.eval()

# 图像预处理
transform = Compose([
    Resize((64, 64)),
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


def allowed_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def predict_image(image_path):
    """预测图像类别"""
    try:
        image = Image.open(image_path).convert('RGB')
        image_tensor = transform(image).unsqueeze(0).to(device)

        with torch.no_grad():
            output = model(image_tensor)
            _, predicted = torch.max(output.data, 1)
            prediction = predicted.item()

        return str(prediction)  # 直接返回预测的数字

    except Exception as e:
        print(f"预测错误: {str(e)}")
        return None


@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        # 检查是否有文件上传
        if 'file' not in request.files:
            return render_template('index.html', error="没有选择文件")

        file = request.files['file']

        if file.filename == '':
            return render_template('index.html', error="没有选择文件")

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)

            # 进行预测
            prediction = predict_image(filepath)

            if prediction is not None:
                return render_template('index.html',
                                       prediction=prediction,
                                       image_path=filepath)
            else:
                return render_template('index.html',
                                       error="预测失败，请重试")

    return render_template('index.html')


@app.route('/predict', methods=['POST'])
def predict_api():
    """API接口，返回JSON格式结果"""
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400

    file = request.files['file']

    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    if file and allowed_file(file.filename):
        # 直接在内存中处理，不保存文件
        image = Image.open(io.BytesIO(file.read())).convert('RGB')
        image_tensor = transform(image).unsqueeze(0).to(device)

        with torch.no_grad():
            output = model(image_tensor)
            _, predicted = torch.max(output.data, 1)
            prediction = predicted.item()

        return jsonify({
            'prediction': int(prediction),  # 返回数字
            'confidence': torch.softmax(output, dim=1)[0][prediction].item()
        })

    return jsonify({'error': 'Invalid file type'}), 400


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)